Summary

Dev Engineer Review

• Added no_repeat_ngram_size support to TorchSampler, aligning with C++ banRepeatNgram and Hugging Face behavior:
  • Prevents sampling tokens that would recreate an existing n-gram in the generated sequence including prompt tokens.
  • Enforced per beam.
  • When n == 1, bans all previously seen tokens (unigram repeat prevention).
  • Feature is disabled for None or 0.
• Refactored bad-word and no-repeat-ngram logic to use a shared suffix-rule ban framework:
  • Unified suffix-rule ban collection with a single shared flush via _apply_token_bans.
  • Added memoized token/context access to reduce host overhead.
  • Introduced _TokenBans to accumulate ban targets across unconditional, overlap-conditioned, and device-token-column cases.
  • Implemented incremental per-request n-gram state indexing via _NgramIndexCache for efficient updates.
  • Reused the existing neg_inf tensor for overlap conditional penalties to avoid extra allocations.
  • Updated _process_requests to compute per-request ngram_sizes and apply both features via shared collection/flush when overlap scheduler tracking is enabled.
  • Improved cache lifecycle by storing _NgramIndexCache in WeakKeyDictionary[LlmRequest, _NgramIndexCache] so state is reclaimed when requests are freed (including rollback scenarios).
• Updated request plumbing:
  • executor_request_to_llm_request now reads executor_request.sampling_config.no_repeat_ngram_size, normalizes falsy/disabled values to None, and sets llm_request.py_no_repeat_ngram_size.
• Documentation/API clarity:
  • Updated docs/source/features/sampling.md to document prompt-inclusive n-gram banning and that None/0 disables the restriction.
  • Updated SamplingParams docstring to explicitly describe the prompt-inclusive behavior and the default (None).

QA Engineer Review

• Test code changes (CUDA unit tests) in tests/unittest/_torch/sampler/test_torch_sampler.py:
  • Modified TestApplyBadWords to use a new sampler helper for both fresh and stale-host paths.
  • Added CUDA-only TestApplyNoRepeatNgram with a MockLlmRequest and helpers to validate banned-logit columns and correct index/cache behavior.
  • Added 18 test cases covering: matching behavior, disabled requests, speculative decoding, multi-beam layouts, overlap handling, duplicate bans (no NaNs), incremental cache updates, cache growth/rollback invalidation, and stale-host execution paths (device hit/miss, mixed stale/fresh batches, window ending at device token).
• Integration/manual coverage in tests/integration/test_lists/, test-db/, or qa/: not verifiable from the provided context (no test-list changes available).
  • Verdict: needs follow-up.

Description

This PR does two things:

1. Implements no_repeat_ngram_size for the PyTorch TorchSampler. The
   parameter was previously honored only by the C++ executor path and silently
   ignored on the PyTorch backend.
2. Refactors all token-ban handling (bad words, no-repeat ngram, and the
   min-length EOS suppression) out of the large sampler.py into a dedicated
   token_ban submodule with a small, testable class hierarchy.

no_repeat_ngram_size semantics

A token is excluded from sampling if it would recreate an n-gram already
present in the sequence (prompt included) — same semantics as the C++
banRepeatNgram kernel and HF's NoRepeatNGramLogitsProcessor. n == 1 bans
every token already present; the restriction is per-beam. None or 0
disables it. Docs (docs/source/features/sampling.md) and the SamplingParams
docstring are updated accordingly.

Token-ban submodule (token_ban.py)

Bad words, no-repeat ngram, and min-length EOS suppression are all the same
shape of operation — mask specific logits to -inf. They now share one module:

• TokenBans — host-side accumulator of (row, col) index lists. Every
  feature appends into a single instance, so all bans reach the device with one
  H2D transfer + one index_put_ per category (instead of one set per feature).
• TokenBanHandler (ABC) — holds the shared feature logic: the bad-words /
  no-repeat-ngram rule generators, the suffix-rule matcher, the incremental
  per-request n-gram index, min-length collection, and the unconditional flush.
• Two variants, selected once at construction from whether the overlap
  scheduler is enabled (a fixed config choice, not per-request):
  • SynchronousTokenBanHandler — overlap off; the host history is always
    complete, so every ban is unconditional.
  • OverlappedTokenBanHandler — overlap on; a batch mixes fresh requests
    (history complete) with stale ones (missing the previous step's token,
    still on the device). Stale requests produce conditional bans resolved on
    the device at apply time (torch.where on the pending token) without a
    device-to-host sync.

TorchSampler now just owns a handler and calls
generate_ban_list() / apply_ban_list(); it computes the per-request stale
flags since it owns the pending-step / draft-batch context.

Performance

The n-gram matcher keeps an incremental per-request index rather than
rescanning the history each step: a host-side token mirror is grown via the
scalar get_num_tokens / get_last_tokens accessors, avoiding the ~40 µs full
get_tokens() copy, so the steady-state per-step cost is O(new tokens).

Isolated sampler benchmark (H200, bs=32, 2048-token context per request, n=3,
p50 sampler-beyond-forward):

The submodule refactor and folding in min-length introduced no measurable
change vs the pre-refactor numbers.

Overlap + speculation / beam search

As with bad words, when the overlap scheduler is combined with speculative
decoding or beam search the missing host tokens cannot be reconstructed with
the single-token device check, so bans may be enforced one step late (warned
once). The exact stale path only runs for single-step, single-beam requests.

Test Coverage

tests/unittest/_torch/sampler/test_token_ban.py (new file) covers all three
token-ban paths against the handlers directly:

• Bad words — single-/multi-token words, prefix hit/miss, prompt-internal
  prefixes, and the full stale-host overlap matrix.
• No-repeat ngram — bigram/trigram hit & miss, multiple matches, n == 1,
  sequences shorter than n, disabled (None and 0), multi-step and
  multi-beam (per-beam history) layouts, the stale overlap matrix (device
  hit/miss, window ending at the device token, mixed stale+fresh, no NaN on
  duplicate bans), incremental cache growth, and cache rebuild on rollback.
• Min length — EOS suppressed below the minimum (generated length excludes
  the prompt), no ban once reached, invalid end_id skipped, multi-step boundary.

Tests assert that logits outside the banned set are left unchanged (starting
from distinct non-zero logits), not just that the banned cells become -inf.

All pass on H200. Additional randomized fuzz (matched against a brute-force
reference and against the "fresh path on the completed sequence" for the stale
path) was run during development.

PR Checklist

• PR title follows the [JIRA/NVBUG/None][type] description format
• Documentation updated (docs/source/features/sampling.md, SamplingParams docstring)
• Unit tests added and passing
• pre-commit hooks pass

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.